#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Time        : 2021/11/19 14:06
@Author      : Albert Darren
@Contact     : 2563491540@qq.com
@File        : grab_cut.py
@Version     : Version 1.0.0
@Description : TODO
@Created By  : PyCharm
"""
import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt

img = cv.imread("../experiment_fig/whale.png")
img1 = np.copy(img)
drawing = False
ix, iy = -1, -1
nx, ny = 0, 0

def draw(event, x, y, flag, param):
    global drawing, ix, iy, nx, ny
    if event == cv.EVENT_LBUTTONDOWN:
        drawing = True
        ix, iy = x, y
        print(x, y)
    if event == cv.EVENT_LBUTTONUP:
        cv.rectangle(img1, (ix, iy), (x, y), (0, 255, 0), 2)
        drawing = False
        nx, ny = x, y
        print(x, y, "/", nx, ny)


for i in range(5):
    while 1:
        cv.imshow("image", img1)
        k = cv.waitKey(1) & 0xFF
        if k == 27:
            break
        cv.setMouseCallback("image", draw)
    cv.destroyAllWindows()
    mask = np.zeros(img.shape[:2], np.uint8)
    bgdModel = np.zeros((1, 65), np.float64)
    fgdModel = np.zeros((1, 65), np.float64)
    rect = (ix, iy, nx, ny)
    cv.grabCut(img, mask, rect, bgdModel, fgdModel, 5, cv.GC_INIT_WITH_RECT)
    mask2 = np.where((mask == 2) | (mask == 0), 0, 1).astype("uint8")
    img = img * mask2[:, :, np.newaxis]
    plt.imshow(img)
    plt.show()
